import os
from transformers.models.bert.tokenization_bert import *
import torch
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
from diffusers.utils import load_image

os.environ["http_proxy"] = "http://192.168.3.116:7890/"
os.environ["https_proxy"] = "http://192.168.3.116:7890/"

import torch
from diffusers import StableDiffusionPipeline
from transformers import AutoTokenizer, CLIPTextModel

# tokenizer_id = "lyua1225/clip-huge-zh-75k-steps-bs4096"
# # tokenizer_id = "Midu/chinese-style-stable-diffusion-2-v0.1"
# sd2_id = "Midu/chinese-style-stable-diffusion-2-v0.1"
# # text_encoder = CLIPTextModel.from_pretrained(tokenizer_id).half()
# tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, trust_remote_code=True)
# pipe = StableDiffusionPipeline.from_pretrained(sd2_id, torch_dtype=torch.float16, tokenizer=tokenizer)
# # pipe = StableDiffusionPipeline.from_pretrained(sd2_id, torch_dtype=torch.float16)
# pipe.to("cuda")
#
# image = pipe("赛博朋克风格的城市街道，8K分辨率，CG渲染", guidance_scale=10, num_inference_steps=20).images[0]
# image.save("cyberpunk.jpeg")

image_path="data/12d/00000.png"
image_path="data/13l/00000.png"
# image_path="data/image/3.png"
init_image = load_image(image_path)

tokenizer_id = "lyua1225/clip-huge-zh-75k-steps-bs4096"
# tokenizer_id = "Midu/chinese-style-stable-diffusion-2-v0.1"
sd2_id = "Midu/chinese-style-stable-diffusion-2-v0.1"
# text_encoder = CLIPTextModel.from_pretrained(tokenizer_id).half()
tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, trust_remote_code=True)
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(sd2_id, torch_dtype=torch.float16, tokenizer=tokenizer)
# pipe = StableDiffusionPipeline.from_pretrained(sd2_id, torch_dtype=torch.float16)
pipe.to("cuda")

# image = pipe("赛博朋克风格的城市街道，8K分辨率，CG渲染", guidance_scale=10, num_inference_steps=20).images[0]
# strength 重绘幅度
image = pipe("一只小柴犬，2d动漫风格，8K分辨率，CG渲染", image=init_image, guidance_scale=7.5,strength=0.6).images[0]
image.save("chinese-style.jpeg")

